#include "DepthRenderTarget.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

DepthRenderTarget::BuilderDetails::BuilderDetails()
{
	isMultiple = false;

	format = GVulkanRHI->GetDepthFormat();
	width = GRenderState->GetWidth();
	height = GRenderState->GetHeight();
	arrayLayers = 1;
	// samples = GVulkanRHI->GetMsaaSamples();
	samples = VK_SAMPLE_COUNT_1_BIT;
	usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
	sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	// aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT;
	aspectMask = 0;

	loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
	storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
	// If you want to read this depth texture on other renderpass , set finalLayout to VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL
	finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

	clearDepthStencilValue = { 1.f, 0 };
}

DepthRenderTarget::Builder::Builder() noexcept = default;
DepthRenderTarget::Builder::~Builder() noexcept = default;
DepthRenderTarget::Builder::Builder(DepthRenderTarget::Builder const& rhs) noexcept = default;
DepthRenderTarget::Builder::Builder(DepthRenderTarget::Builder && rhs) noexcept = default;
DepthRenderTarget::Builder& DepthRenderTarget::Builder::operator=(DepthRenderTarget::Builder const& rhs) noexcept = default;
DepthRenderTarget::Builder& DepthRenderTarget::Builder::operator=(DepthRenderTarget::Builder && rhs) noexcept = default;

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetMultiple(bool isMultiple) noexcept
{
	mImpl->isMultiple = isMultiple;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetFormat(VkFormat format) noexcept
{
	mImpl->format = format;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetWidth(uint32 width) noexcept
{
	mImpl->width = width;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetHeight(uint32 height) noexcept
{
	mImpl->height = height;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetArrayLayers(uint32 arrayLayers) noexcept
{
	mImpl->arrayLayers = arrayLayers;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetSamples(VkSampleCountFlagBits samples) noexcept
{
	assert(samples <= GVulkanRHI->GetMsaaSamples());
	mImpl->samples = samples;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetUsage(VkImageUsageFlags usage) noexcept
{
	// If use for msaa to make the resolve texture, need VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT
	// If use for input attachment, need VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT
	// Depth texture need to contain VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT
	assert(usage & VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT);
	mImpl->usage = usage;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetAspectMask(VkImageAspectFlags aspectMask) noexcept
{
	mImpl->aspectMask = aspectMask;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetLoadOp(VkAttachmentLoadOp loadOp) noexcept
{
	mImpl->loadOp = loadOp;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetStoreOp(VkAttachmentStoreOp storeOp) noexcept
{
	mImpl->storeOp = storeOp;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetStencilLoadOp(VkAttachmentLoadOp stencilLoadOp) noexcept
{
	mImpl->stencilLoadOp = stencilLoadOp;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetStencilStoreOp(VkAttachmentStoreOp stencilStoreOp) noexcept
{
	mImpl->stencilStoreOp = stencilStoreOp;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetInitialLayout(VkImageLayout initialLayout) noexcept
{
	mImpl->initialLayout = initialLayout;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetFinalLayout(VkImageLayout finalLayout) noexcept
{
	mImpl->finalLayout = finalLayout;
	return *this;
}

DepthRenderTarget::Builder& DepthRenderTarget::Builder::SetClearDepthStencilValue(VkClearDepthStencilValue clearDepthStencilValue) noexcept
{
	mImpl->clearDepthStencilValue = clearDepthStencilValue;
	return *this;
}

RefCountPtr<DepthRenderTarget> DepthRenderTarget::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<DepthRenderTarget>(NewObject<DepthRenderTarget>(*mImpl));
}

RefCountPtr<DepthRenderTarget> DepthRenderTarget::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	DepthRenderTarget* depthRenderTarget = NewObject<DepthRenderTarget>(*mImpl);
	process->AddTexture(name, depthRenderTarget);
	return RefCountPtr<DepthRenderTarget>(depthRenderTarget);
}

void DepthRenderTarget::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveTexture(mName, this);
		}
		DeleteObject<DepthRenderTarget>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveTexture(mName, this);
			}
			DeleteObject<DepthRenderTarget>(this);
			});
	}
}

DepthRenderTarget::DepthRenderTarget(BuilderDetails & builderDetails)
{
	mIsMultiple = builderDetails.isMultiple;
	mIsRenderTarget = true;
	mRenderTargetType = RenderTargetType::Depth;

	mFlags = 0;
	mImageType = VK_IMAGE_TYPE_2D;
	mFormat = builderDetails.format;
	mExtent = { builderDetails.width, builderDetails.height, 1 };
	mUsage = builderDetails.usage;
	mMipLevels = 1;
	mArrayLayers = builderDetails.arrayLayers;
	mSamples = builderDetails.samples;
	mTiling = VK_IMAGE_TILING_OPTIMAL;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// Check if the queueFamilys is same
		bool isDifferent = false;
		const uint32 standQueueFamily = GVulkanRHI->GetQueueFamily(mQueueFamilys[0]);
		for (uint32 i = 1; i < mQueueFamilys.size(); ++i)
		{
			if (GVulkanRHI->GetQueueFamily(mQueueFamilys[i]) != standQueueFamily)
				isDifferent = true;
		}

		if (isDifferent)
			mSharingMode = VK_SHARING_MODE_CONCURRENT; // multi queue
		else
			mSharingMode = VK_SHARING_MODE_EXCLUSIVE;
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mImageLayouts.resize(GetFrameCount(), VK_IMAGE_LAYOUT_UNDEFINED);
	mMemoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;

	mViewType = mArrayLayers == 1 ? VK_IMAGE_VIEW_TYPE_2D : VK_IMAGE_VIEW_TYPE_2D_ARRAY;
	mComponents = { VK_COMPONENT_SWIZZLE_R , VK_COMPONENT_SWIZZLE_G , VK_COMPONENT_SWIZZLE_B , VK_COMPONENT_SWIZZLE_A };

	
	if (builderDetails.aspectMask == 0)
	{
		if (HasDepth())
			builderDetails.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
		if (HasStencil())
			builderDetails.aspectMask = builderDetails.aspectMask | VK_IMAGE_ASPECT_STENCIL_BIT;
	}
	else
	{
		// When using a depth/stencil image in a descriptor set, please only set either VK_IMAGE_ASPECT_DEPTH_BIT or VK_IMAGE_ASPECT_STENCIL_BIT
		if (mUsage & VK_IMAGE_USAGE_SAMPLED_BIT)
		{
			if (builderDetails.aspectMask & VK_IMAGE_ASPECT_DEPTH_BIT)
				assert(!(builderDetails.aspectMask & VK_IMAGE_ASPECT_STENCIL_BIT));
			if(builderDetails.aspectMask & VK_IMAGE_ASPECT_STENCIL_BIT)
				assert(!(builderDetails.aspectMask & VK_IMAGE_ASPECT_DEPTH_BIT));
		}
	}

	assert(builderDetails.aspectMask > 0);
	mSubresourceRange = { builderDetails.aspectMask, 0, 1, 0 , mArrayLayers };

	mAttachmentDescription = {};
	mAttachmentDescription.format = mFormat;
	mAttachmentDescription.samples = mSamples;
	mAttachmentDescription.loadOp = builderDetails.loadOp;
	mAttachmentDescription.storeOp = builderDetails.storeOp;
	mAttachmentDescription.stencilLoadOp = builderDetails.stencilLoadOp;
	mAttachmentDescription.stencilStoreOp = builderDetails.stencilStoreOp;
	mAttachmentDescription.initialLayout = builderDetails.initialLayout;
	mAttachmentDescription.finalLayout = builderDetails.finalLayout;

	mClearValue = {};
	mClearValue.depthStencil = builderDetails.clearDepthStencilValue;

	mImages.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageMemorys.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageViews.resize(GetFrameCount(), VK_NULL_HANDLE);

	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateImage(mImages[i], mImageMemorys[i], mFlags, mExtent, mUsage, mFormat, mMipLevels, mMemoryProperty, mSamples, mImageType, mImageLayouts[i], mArrayLayers, mTiling, mSharingMode, mQueueFamilys);

		GVulkanRHI->CreateImageView(mImageViews[i], mImages[i], mFormat, mViewType, mSubresourceRange, mComponents);

		// GVulkanRHI->TransferLayout(commandBuffer, mImages[i], mImageLayouts[i], VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT);
		// mImageLayouts[i] = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
	}

	GRenderState->FlushSingleCommand(commandBuffer);
}

bool DepthRenderTarget::HasDepth() const
{
	const static std::vector<VkFormat> formats =
	{
		VK_FORMAT_D16_UNORM,
		VK_FORMAT_X8_D24_UNORM_PACK32,
		VK_FORMAT_D32_SFLOAT,
		VK_FORMAT_D16_UNORM_S8_UINT,
		VK_FORMAT_D24_UNORM_S8_UINT,
		VK_FORMAT_D32_SFLOAT_S8_UINT
	};
	return std::find(formats.begin(), formats.end(), mFormat) != std::end(formats);
}

bool DepthRenderTarget::HasStencil() const
{
	const static std::vector<VkFormat> formats =
	{
		VK_FORMAT_S8_UINT,
		VK_FORMAT_D16_UNORM_S8_UINT,
		VK_FORMAT_D24_UNORM_S8_UINT,
		VK_FORMAT_D32_SFLOAT_S8_UINT
	};
	return std::find(formats.begin(), formats.end(), mFormat) != std::end(formats);
}